package net.cyue.loader;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 增强版带缓存功能的ClassLoader，适用于各种类加载场景
 * 继承自基础ClassLoader，实现了类缓存和并发控制
 */
public class CacheClassLoader extends ClassLoader {
    // 类缓存，使用ConcurrentHashMap保证线程安全，初始容量和负载因子优化
    private final Map<String, LoadClassResult> classCache = new ConcurrentHashMap<>(2048, 0.75f);

    // 用于防止并发加载同一类的锁，使用ConcurrentHashMap避免同步块
    private final Map<String, Lock> locks = new ConcurrentHashMap<>(128);

    // 缓存未找到类的常量实例
    private static final LoadClassResult NOT_FOUND = new LoadClassResult(null, null);

    // 缓冲区重用，减少对象创建
    private static final ThreadLocal<byte[]> BUFFER_CACHE = ThreadLocal.withInitial(() -> new byte[8192]);

    /**
     * 构造函数，使用默认的父类加载器
     */
    public CacheClassLoader() {
        super();
    }

    /**
     * 构造函数，指定父类加载器
     * @param parent 父类加载器
     */
    public CacheClassLoader(ClassLoader parent) {
        super(parent);
    }

    @Override
    public Class<?> loadClass(String name) throws ClassNotFoundException {
        return loadClass(name, false);
    }

    @Override
    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
        // 对于Java核心类，直接使用父类加载器加载，不缓存
        if (name.startsWith("java.") || name.startsWith("javax.") || name.startsWith("sun.")) {
            return super.loadClass(name, resolve);
        }
        return loadClassWithCache(name, resolve);
    }

    /**
     * 带缓存和并发控制的类加载方法
     */
    private Class<?> loadClassWithCache(String name, boolean resolve) throws ClassNotFoundException {
        // 先从缓存获取
        LoadClassResult result = classCache.get(name);
        if (result != null) {
            return processCachedResult(name, result, resolve);
        }

        // 缓存未命中，获取或创建对应类的锁
        Lock lock = locks.computeIfAbsent(name, k -> new ReentrantLock());
        boolean locked = false;
        try {
            // 尝试获取锁，设置超时防止死锁风险
            locked = lock.tryLock();
            if (!locked) {
                // 未能获取锁，可能有其他线程在加载，短暂等待后再次检查缓存
                Thread.yield();
                result = classCache.get(name);
                if (result != null) {
                    return processCachedResult(name, result, resolve);
                }
                // 再次尝试获取锁，这次阻塞等待
                lock.lock();
                locked = true;
            }

            // 双重检查缓存
            result = classCache.get(name);
            if (result != null) {
                return processCachedResult(name, result, resolve);
            }

            // 执行实际加载
            return performLoadClass(name, resolve);
        } finally {
            if (locked) {
                lock.unlock();
            }
        }
    }

    /**
     * 处理缓存结果
     */
    private Class<?> processCachedResult(String name, LoadClassResult result, boolean resolve)
            throws ClassNotFoundException {
        if (result.isSuccess()) {
            Class<?> clazz = result.getClazz();
            if (resolve) {
                resolveClass(clazz);
            }
            return clazz;
        } else if (result.getException() instanceof ClassNotFoundException) {
            throw (ClassNotFoundException) result.getException();
        } else if (result.getException() != null) {
            throw new ClassNotFoundException("Class loading failed for: " + name, result.getException());
        } else {
            // 处理NOT_FOUND情况
            throw new ClassNotFoundException("Class not found in cache: " + name);
        }
    }

    /**
     * 执行实际的类加载操作
     */
    private Class<?> performLoadClass(String name, boolean resolve) throws ClassNotFoundException {
        try {
            // 先检查是否已加载
            Class<?> clazz = findLoadedClass(name);
            if (clazz == null) {
                // 尝试用父类加载器加载
                try {
                    if (getParent() != null) {
                        clazz = getParent().loadClass(name);
                    }
                } catch (ClassNotFoundException e) {
                    // 父类加载器未找到，继续尝试自己加载
                }

                // 父类未加载到，尝试自己查找并加载
                if (clazz == null) {
                    byte[] classData = loadClassData(name);
                    if (classData == null) {
                        classCache.put(name, NOT_FOUND);
                        throw new ClassNotFoundException("Class data not found: " + name);
                    }
                    clazz = defineClass(name, classData, 0, classData.length);
                }
            }

            if (clazz == null) {
                classCache.put(name, NOT_FOUND);
                throw new ClassNotFoundException("Class not found: " + name);
            }

            // 存入缓存
            LoadClassResult successResult = new LoadClassResult(clazz, null);
            classCache.put(name, successResult);

            // 处理类解析
            if (resolve) {
                resolveClass(clazz);
            }

            return clazz;
        } catch (ClassNotFoundException e) {
            // 缓存类未找到异常
            classCache.put(name, new LoadClassResult(null, e));
            throw e;
        } catch (LinkageError e) {
            // 缓存链接错误（如类版本不兼容等）
            classCache.put(name, new LoadClassResult(null, e));
            throw new ClassNotFoundException("Linkage error loading class: " + name, e);
        } catch (SecurityException e) {
            // 缓存安全异常
            classCache.put(name, new LoadClassResult(null, e));
            throw new ClassNotFoundException("Security error loading class: " + name, e);
        } catch (IOException e) {
            // 缓存IO异常
            classCache.put(name, new LoadClassResult(null, e));
            throw new ClassNotFoundException("IO error loading class: " + name, e);
        }
    }

    /**
     * 加载类数据，使用NIO优化IO操作，提高读取效率
     * @param className 类名
     * @return 类的字节数组
     * @throws IOException 加载过程中发生IO异常
     */
    protected byte[] loadClassData(String className) throws IOException {
        // 将类名转换为资源路径
        String resourceName = className.replace('.', '/') + ".class";

        // 尝试从类路径加载
        try (InputStream is = getResourceAsStream(resourceName)) {
            if (is == null) {
                return null;
            }

            // 使用NIO通道提高读取效率
            try (ReadableByteChannel channel = Channels.newChannel(is)) {
                ByteArrayOutputStream baos = new ByteArrayOutputStream(4096); // 预设初始容量
                ByteBuffer buffer = ByteBuffer.wrap(BUFFER_CACHE.get());

                int bytesRead;
                while ((bytesRead = channel.read(buffer)) != -1) {
                    buffer.flip();
                    byte[] bytes = new byte[bytesRead];
                    buffer.get(bytes);
                    baos.write(bytes);
                    buffer.clear();
                }

                return baos.toByteArray();
            }
        }
    }

    /**
     * 清除指定类的缓存，同时清除对应的锁以释放资源
     */
    public void clearCache(String className) {
        if (className == null) {
            throw new IllegalArgumentException("Class name cannot be null");
        }
        classCache.remove(className);
        // 可以选择是否移除锁，这里保留锁以避免频繁创建
        // locks.remove(className);
    }

    /**
     * 清除所有类缓存
     */
    public void clearAllCache() {
        classCache.clear();
        // 保留锁以避免频繁创建
        // locks.clear();
    }

    /**
     * 获取缓存大小
     */
    public int getCacheSize() {
        return classCache.size();
    }

    /**
     * 检查类是否在缓存中
     */
    public boolean isCached(String className) {
        return className != null && classCache.containsKey(className);
    }

    /**
     * 类加载结果的封装类，使用final修饰提高性能
     */
    private static final class LoadClassResult {
        private final Class<?> clazz;
        private final Throwable exception;

        public LoadClassResult(Class<?> clazz, Throwable exception) {
            this.clazz = clazz;
            this.exception = exception;
        }

        public Class<?> getClazz() {
            return clazz;
        }

        public Throwable getException() {
            return exception;
        }

        /**
         * 判断加载是否成功
         */
        public boolean isSuccess() {
            return clazz != null && exception == null;
        }
    }
}
